package com.soup.lab.executor.pool2;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class Node<K, E extends KeyPickRunnable<K>> implements Runnable {

    /**
     * 0-初始化
     * 1-执行中
     * 2-已完成
     * 3-已取消
     */
    private final AtomicInteger state = new AtomicInteger();
    /**
     * 任务队列
     */
    private final Queue<E> tasks;
    /**
     * 记录当前参与键及其数量
     */
    private final Map<K, Integer> keyCountMap;

    /**
     * 读写锁
     */
    private final ReentrantReadWriteLock readWriteLock = new ReentrantReadWriteLock();
    private final ReentrantReadWriteLock.WriteLock writeLock = readWriteLock.writeLock();
    private final ReentrantReadWriteLock.ReadLock readLock = readWriteLock.readLock();

    public Node(E firstTask) {
        this.tasks = Lists.newLinkedList();
        this.tasks.offer(firstTask);
        this.keyCountMap = Maps.newHashMap();
    }

    @Override
    public void run() {
        if(!state.compareAndSet(INITIAL, RUNNING)) {
            // 已取消
            return;
        }
        E task = null;
        while ((task = next(task)) != null) {
            task.run();
        }
    }

    private E next(E curr) {
        writeLock.lock();
        try {
            if(isAlive()) {
                if(curr != null) {
                    handleCount(curr, -1);
                }
                E task = tasks.poll();
                if(task == null) {
                    state.set(FINISHED);
                }
                return task;
            }
        } finally {
            writeLock.unlock();
        }
        return null;
    }

    private void handleCount(E runnable, int count) {
        for (K key : runnable.getKeys()) {
            keyCountMap.compute(key, (k,v) -> v == null ? count : v+count == 0 ? null : v+count);
        }
    }

    public boolean add(E task) {
        writeLock.lock();
        try {
           if(isAlive()) {
               handleCount(task, 1);
               return tasks.offer(task);
           }
        } finally {
            writeLock.unlock();
        }
        return false;
    }

    public List<E> cancel() {
        writeLock.lock();
        try {
            state.set(CANCELED);
            return Lists.newArrayList(tasks);
        } finally {
            writeLock.unlock();
        }
    }

    public boolean isAlive() {
        return state.get() < FINISHED;
    }

    public Collection<K> waitingKeySet() {
        readLock.lock();
        try {
            return Sets.newHashSet(keyCountMap.keySet());
        } finally {
            readLock.unlock();
        }
    }

    private static final int INITIAL = 0;
    private static final int RUNNING = 1;
    private static final int FINISHED = 2;
    private static final int CANCELED = 3;

}
